import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tensorflow.keras.datasets import mnist

def save_mnist_as_images(output_dir='mnist_images'):
    """
    将MNIST数据集中的所有图像解压为图片文件
    
    参数:
        output_dir: 输出目录路径(默认'mnist_images')
    """
    # 加载MNIST数据集
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    
    # 创建主输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建训练集和测试集子目录
    train_dir = os.path.join(output_dir, 'train')
    test_dir = os.path.join(output_dir, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)
    
    # 为每个类别(0-9)创建子目录
    for i in range(10):
        os.makedirs(os.path.join(train_dir, str(i)), exist_ok=True)
        os.makedirs(os.path.join(test_dir, str(i)), exist_ok=True)
    
    # 保存训练集图片
    print("正在保存训练集图片...")
    for i in range(len(x_train)):
        img = Image.fromarray(x_train[i])
        label = y_train[i]
        img.save(os.path.join(train_dir, str(label), f'train_{i:05d}.png'))
    
    # 保存测试集图片
    print("正在保存测试集图片...")
    for i in range(len(x_test)):
        img = Image.fromarray(x_test[i])
        label = y_test[i]
        img.save(os.path.join(test_dir, str(label), f'test_{i:05d}.png'))
    
    print(f"所有MNIST图片已保存到: {os.path.abspath(output_dir)}")
    print(f"训练集: {len(x_train)}张图片 (按标签分类在0-9子目录中)")
    print(f"测试集: {len(x_test)}张图片 (按标签分类在0-9子目录中)")

if __name__ == '__main__':
    save_mnist_as_images()